#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据库初始化脚本
确保数据库和表结构正确创建
"""

import os
import sys
import pymysql
import logging
from pathlib import Path
from config import DB_CONFIG

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

def create_database_if_not_exists():
    """创建数据库（如果不存在）"""
    try:
        # 连接MySQL服务器（不指定数据库）
        conn_config = DB_CONFIG.copy()
        database_name = conn_config.pop('database')
        
        logger.info(f"连接到MySQL服务器: {conn_config['host']}:{conn_config['port']}")
        conn = pymysql.connect(**conn_config)
        cursor = conn.cursor()
        
        # 创建数据库
        cursor.execute(f"CREATE DATABASE IF NOT EXISTS `{database_name}` CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci")
        logger.info(f"✅ 数据库 '{database_name}' 已确保存在")
        
        cursor.close()
        conn.close()
        return True
        
    except pymysql.Error as e:
        logger.error(f"❌ 创建数据库失败: {e}")
        return False

def init_tables():
    """初始化数据库表结构"""
    try:
        logger.info("连接到数据库...")
        conn = pymysql.connect(**DB_CONFIG)
        cursor = conn.cursor()
        
        # 创建messages表
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS `messages` (
            `id` INT AUTO_INCREMENT PRIMARY KEY,
            `conversation_id` VARCHAR(255) NOT NULL,
            `sender_type` VARCHAR(50) NOT NULL,
            `content` TEXT,
            `timestamp` DATETIME NOT NULL,
            `created_at` TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            INDEX idx_conversation_id (conversation_id),
            INDEX idx_sender_type (sender_type),
            INDEX idx_timestamp (timestamp)
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
        """)
        logger.info("✅ messages表已创建/确认")
        
        # 创建conversations表
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS `conversations` (
            `id` INT AUTO_INCREMENT PRIMARY KEY,
            `conversation_id` VARCHAR(255) NOT NULL UNIQUE,
            `candidate_name` VARCHAR(255),
            `hr_name` VARCHAR(255),
            `start_time` DATETIME,
            `last_message_time` DATETIME,
            `message_count` INT DEFAULT 0,
            `status` VARCHAR(50),
            `created_at` TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            INDEX idx_conversation_id (conversation_id),
            INDEX idx_status (status)
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
        """)
        logger.info("✅ conversations表已创建/确认")
        
        # 创建evaluations表
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS `evaluations` (
            `id` INT AUTO_INCREMENT PRIMARY KEY,
            `conversation_id` VARCHAR(255) NOT NULL,
            `evaluation_type` VARCHAR(50) NOT NULL,
            `score` INT,
            `reason` TEXT,
            `details` JSON,
            `created_at` TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            INDEX idx_conversation_id (conversation_id),
            INDEX idx_evaluation_type (evaluation_type),
            INDEX idx_score (score)
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
        """)
        logger.info("✅ evaluations表已创建/确认")
        
        # 创建prompts表
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS `prompts` (
            `id` INT AUTO_INCREMENT PRIMARY KEY,
            `prompt_type` VARCHAR(50) NOT NULL UNIQUE,
            `content` TEXT NOT NULL,
            `version` INT DEFAULT 1,
            `updated_at` TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
            INDEX idx_prompt_type (prompt_type)
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
        """)
        logger.info("✅ prompts表已创建/确认")
        
        # 初始化默认提示词
        default_prompts = {
            'salary': '请分析这段对话中HR是否过早询问候选人薪资期望。评分标准：1-3分表示过早询问，4-7分表示时机一般，8-10分表示时机恰当或未询问。',
            'photo': '请分析这段对话中HR是否不当索要候选人照片。评分标准：1-3分表示明显不当，4-7分表示有争议，8-10分表示合理或未索要。',
            'sensitive': '请分析这段对话中是否包含敏感词或不当言论。评分标准：1-3分表示严重违规，4-7分表示轻微违规，8-10分表示无违规。',
            'reply_reasonable': '请分析这段对话中信息对是否合理。评分标准：1-3分表示明显不合理，4-7分表示部分合理，8-10分表示完全合理。'
        }
        
        for prompt_type, content in default_prompts.items():
            cursor.execute("""
                INSERT INTO prompts (prompt_type, content) 
                VALUES (%s, %s)
                ON DUPLICATE KEY UPDATE content = VALUES(content)
            """, (prompt_type, content))
        logger.info("✅ 默认提示词已初始化")
        
        # 创建feedback表
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS `feedback` (
            `id` INT AUTO_INCREMENT PRIMARY KEY,
            `conversation_id` VARCHAR(255) NOT NULL,
            `evaluation_type` VARCHAR(50) NOT NULL,
            `score` INT,
            `suggestion` TEXT,
            `tags` JSON,
            `created_at` TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            INDEX idx_conversation_id (conversation_id),
            INDEX idx_evaluation_type (evaluation_type)
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
        """)
        logger.info("✅ feedback表已创建/确认")
        
        conn.commit()
        cursor.close()
        conn.close()
        
        logger.info("✅ 所有表结构初始化完成")
        return True
        
    except pymysql.Error as e:
        logger.error(f"❌ 初始化表结构失败: {e}")
        return False

def insert_test_data():
    """插入测试数据"""
    try:
        conn = pymysql.connect(**DB_CONFIG)
        cursor = conn.cursor()
        
        # 插入测试对话
        test_conversations = [
            ('test_001', '张三', 'HR小王'),
            ('test_002', '李四', 'HR小李'),
            ('test_003', '王五', 'HR小张'),
        ]
        
        for conv_id, candidate, hr in test_conversations:
            cursor.execute("""
                INSERT INTO conversations (conversation_id, candidate_name, hr_name, start_time, last_message_time, message_count, status)
                VALUES (%s, %s, %s, NOW(), NOW(), 2, 'active')
                ON DUPLICATE KEY UPDATE last_message_time = NOW()
            """, (conv_id, candidate, hr))
            
            # 插入测试消息
            cursor.execute("""
                INSERT INTO messages (conversation_id, sender_type, content, timestamp)
                VALUES (%s, 'hr', '您好，请问您期望的薪资是多少？', NOW())
            """, (conv_id,))
            
            cursor.execute("""
                INSERT INTO messages (conversation_id, sender_type, content, timestamp)
                VALUES (%s, 'candidate', '我期望月薪8000-10000元', NOW())
            """, (conv_id,))
        
        conn.commit()
        logger.info("✅ 测试数据已插入")
        
        cursor.close()
        conn.close()
        return True
        
    except pymysql.Error as e:
        logger.error(f"❌ 插入测试数据失败: {e}")
        return False

def check_database_status():
    """检查数据库状态"""
    try:
        conn = pymysql.connect(**DB_CONFIG)
        cursor = conn.cursor()
        
        # 检查各表的记录数
        tables = ['messages', 'conversations', 'evaluations', 'prompts', 'feedback']
        logger.info("\n数据库状态检查:")
        logger.info("-" * 40)
        
        for table in tables:
            cursor.execute(f"SELECT COUNT(*) as count FROM {table}")
            result = cursor.fetchone()
            count = result[0] if result else 0
            logger.info(f"{table:20} : {count} 条记录")
        
        cursor.close()
        conn.close()
        return True
        
    except pymysql.Error as e:
        logger.error(f"❌ 检查数据库状态失败: {e}")
        return False

def main():
    """主函数"""
    logger.info("="*50)
    logger.info("开始初始化数据库")
    logger.info("="*50)
    
    # 1. 创建数据库
    if not create_database_if_not_exists():
        logger.error("创建数据库失败，退出")
        return False
    
    # 2. 初始化表结构
    if not init_tables():
        logger.error("初始化表结构失败，退出")
        return False
    
    # 3. 插入测试数据（可选）
    response = input("\n是否插入测试数据？(y/n): ")
    if response.lower() == 'y':
        insert_test_data()
    
    # 4. 检查数据库状态
    check_database_status()
    
    logger.info("\n✅ 数据库初始化完成！")
    logger.info("现在可以运行 python start_server.py 启动服务器")
    return True

if __name__ == "__main__":
    main()
